# ======================================================================
# Copyright CERFACS (January 2019)
# Contributor: Adrien Suau (adrien.suau@cerfacs.fr)
#
# This software is governed by the CeCILL-B license under French law and
# abiding  by the  rules of  distribution of free software. You can use,
# modify  and/or  redistribute  the  software  under  the  terms  of the
# CeCILL-B license as circulated by CEA, CNRS and INRIA at the following
# URL "http://www.cecill.info".
#
# As a counterpart to the access to  the source code and rights to copy,
# modify and  redistribute granted  by the  license, users  are provided
# only with a limited warranty and  the software's author, the holder of
# the economic rights,  and the  successive licensors  have only limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using, modifying and/or  developing or reproducing  the
# software by the user in light of its specific status of free software,
# that  may mean  that it  is complicated  to manipulate,  and that also
# therefore  means that  it is reserved for  developers and  experienced
# professionals having in-depth  computer knowledge. Users are therefore
# encouraged  to load and  test  the software's  suitability as  regards
# their  requirements  in  conditions  enabling  the  security  of their
# systems  and/or  data to be  ensured and,  more generally,  to use and
# operate it in the same conditions as regards security.
#
# The fact that you  are presently reading this  means that you have had
# knowledge of the CeCILL-B license and that you accept its terms.
# ======================================================================

"""Module implementing Hamiltonian simulation for 1-sparse fixed point Hamiltonians.

This module contains the functions used to simulate a 1-sparse Hamiltonian that contain
either only real fixed-point entries or only imaginary fixed-point entries.

The circuit returned by these procedures simulate exactly the given Hamiltonian, but
encoding real numbers on a fixed-point representation introduce an error in the
Hamiltonian that will be reflected in the simulation. In other words, the routines
simulate exactly the provided Hamitlonians, but the provided Hamiltonians are not always
exactly the Hamiltonians we wish to simulate because of the fixed-point representation
error.
"""

from qaths.simulation.base._routines import A, G, exp_ZZF2r_t
from qat.lang.AQASM.routines import QRoutine
from qat.lang.AQASM.misc import build_gate
from qat.lang.AQASM.gates import AbstractGate


@build_gate(
    "simulate_float_weighted_hamiltonian",
    [str, int, int, float],
    arity=lambda name, n, r, t: n,
)
def simulate_float_weighted_hamiltonian(
    oracle_name: str, n: int, r: int, time: float
) -> QRoutine:
    r"""Simulate the evolution under an Hamiltonian with at most one real entry in each\
    row/column.

    Let :math:`n` be the size of :math:`\ket{x}`. The returned routine needs
    :math:`n` qubits organised as follow: ::

        |    |x>    |
        |   .    .  |
        |  0    n-1 |


    * x is the starting state of the evolution
      :math:`\left(\ket{x} = \ket{\phi_0}\right)`.

    :param oracle_name: name of the oracle that encodes the Hamiltonian matrix we
        want to simulate.
        The name will be used to create an AbstractGate, that can be linked afterwards.

        The oracle :math:`O` take as input :math:`\ket{x}\ket{0}\ket{0}\ket{0}`
        organised as follow: ::

            |      |x>      |        |m>        |        |v>        |  |s>  |
            |   .   .   .   |   .   .   .   .   |   .   .   .   .   |   .   |
            |  0       n-1  |   n         2*n-1 |  2*n       2n+3r-1| 2n+3r |

        and outputs the quantum state
        :math:`\ket{x}\ket{m(x)}\ket{v(x)}\ket{s(x)}` with:

        #. :math:`\ket{x}` encoding the index of the considered row (all the indices are
           in superposition).
        #. :math:`\ket{m(x)}` encoding the index of the column of the only non-zero
           element in the column of index :math:`x`.
        #. :math:`\ket{v(x)}` encoding the absolute value of the only non-zero entry
           in the column of index :math:`x`.
        #. :math:`\ket{s(x)}` encoding the sign of the only non-zero entry in the column
           of index :math:`x` with the convention

           .. math::

              s(x) = \left\{\begin{array}{lr}
              0 & \text{if } x \geqslant 0 \\
              1 & \text{else} \\
              \end{array}\right.

    :param n: The number of qubits the simulated Hamiltonian acts on.
    :param r: Number of qubits used to represent fixed-point real numbers.
    :param time: The duration of the desired evolution.
    """
    routine = QRoutine()
    # Aliases to make the code more readable.
    x = routine.new_wires(n)
    m = routine.new_wires(n)
    v = routine.new_wires(3 * r)
    p = routine.new_wires(1)
    s = routine.new_wires(1)
    routine.set_ancillae(m)
    routine.set_ancillae(v)
    routine.set_ancillae(p)
    routine.set_ancillae(s)

    oracle = AbstractGate(oracle_name, [], arity=2 * n + 3 * r + 1)

    # First, apply the oracle on the provided qubits
    routine.protected_apply(oracle, x, m, v, s)
    routine.protected_apply(A(n), x, m, p)

    routine.apply(exp_ZZF2r_t(r, time), p, s, v)

    routine.protected_apply(A(n).dag(), x, m, p)
    # Last, apply the inversed oracle on the given qubits
    routine.protected_apply(oracle.dag(), x, m, v, s)

    return routine


@build_gate(
    "simulate_imaginary_float_weighted_hamiltonian",
    [str, int, int, float],
    arity=lambda name, n, r, t: n,
)
def simulate_imaginary_float_weighted_hamiltonian(
    oracle_name: str, n: int, r: int, time: float
) -> QRoutine:
    r"""Simulate the evolution under an Hamiltonian with at most one positive \
    imaginary entry in each row/column.

    Let :math:`n` be the size of :math:`\ket{x}`. The returned routine needs
    :math:`n` qubits organised as follow: ::

        |    |x>    |
        |   .    .  |
        |  0    n-1 |


    * x is the starting state of the evolution
      :math:`\left(\ket{x} = \ket{\phi_0}\right)`.

    :param oracle_name: name of the oracle that encodes the Hamiltonian matrix we
        want to simulate.
        The name will be used to create an AbstractGate, that can be linked afterwards.


        The oracle :math:`O` take as input :math:`\ket{x}\ket{0}\ket{0}\ket{0}`
        organised as follow: ::

            |      |x>      |        |m>        |        |v>        |  |s>  |
            |   .   .   .   |   .   .   .   .   |   .   .   .   .   |   .   |
            |  0       n-1  |   n         2*n-1 |  2*n       2n+3r-1| 2n+3r |

        and outputs the quantum state
        :math:`\ket{x}\ket{m(x)}\ket{v(x)}\ket{s(x)}` with:

        #. :math:`\ket{x}` encoding the index of the considered row (all the indices are
           in superposition).
        #. :math:`\ket{m(x)}` encoding the index of the column of the only non-zero
           element in the column of index :math:`x`.
        #. :math:`\ket{v(x)}` encoding the absolute value of the only non-zero entry
           in the column of index :math:`x`.
        #. :math:`\ket{s(x)}` encoding the sign of the only non-zero entry in the column
           of index :math:`x` with the convention

           .. math::

              s(x) = \left\{\begin{array}{lr}
              0 & \text{if } x \geqslant 0 \\
              1 & \text{else} \\
              \end{array}\right.

    :param n: The number of qubits the simulated Hamiltonian acts on.
    :param r: Number of qubits used to represent fixed-point real numbers.
    :param time: The duration of the desired evolution.
    """
    routine = QRoutine()
    # Aliases to make the code more readable.
    x = routine.new_wires(n)
    m = routine.new_wires(n)
    v = routine.new_wires(3 * r)
    p = routine.new_wires(1)
    s = routine.new_wires(1)
    routine.set_ancillae(m)
    routine.set_ancillae(v)
    routine.set_ancillae(p)
    routine.set_ancillae(s)

    oracle = AbstractGate(oracle_name, [], arity=2 * n + 3 * r + 1)

    # Prepare the quantum state
    routine.protected_apply(oracle, x, m, v, s)
    routine.protected_apply(A(n), x, m, p)
    routine.protected_apply(G, s)

    # Perform the simulation
    routine.apply(exp_ZZF2r_t(r, time), p, s, v)

    # Inverse the preparatory step.
    routine.protected_apply(G.dag(), s)
    routine.protected_apply(A(n).dag(), x, m, p)
    routine.protected_apply(oracle.dag(), x, m, v, s)

    return routine
